from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
import os
import sqlite3
import torch
from torchvision import models, transforms
from PIL import Image
import numpy as np
from flask import Flask, request, jsonify, send_from_directory
import json

# 初始化模型
model = models.resnet50(pretrained=True)
model.eval()

# 定义文件夹路径
UPLOAD_FOLDER = 'upload_folder'
TMP_FOLDER = 'temp'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)
if not os.path.exists(TMP_FOLDER):
    os.makedirs(TMP_FOLDER)

# 数据库初始化函数
def init_db():
    conn = sqlite3.connect('images.db')
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS images (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            path TEXT NOT NULL,
            feature BLOB NOT NULL
        )
    ''')
    conn.commit()
    conn.close()

# 特征提取
def extract_features(image_path):
    img = Image.open(image_path).convert('RGB')
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    img = preprocess(img).unsqueeze(0)
    with torch.no_grad():
        features = model(img).numpy().flatten()
    return features

# 数据库操作
def insert_image(path, feature):
    conn = sqlite3.connect('images.db')
    c = conn.cursor()
    print(path)
    c.execute("INSERT INTO images (path, feature) VALUES (?, ?)", (path, feature.tobytes()))
    conn.commit()
    conn.close()

def search_similar_images(query_feature, top_k=5):
    conn = sqlite3.connect('images.db')
    c = conn.cursor()
    c.execute("SELECT path, feature FROM images")
    results = []
    for row in c.fetchall():
        path, feature_blob = row
        feature = np.frombuffer(feature_blob, dtype=np.float32)
        distance = np.linalg.norm(query_feature - feature)
        # results.append((path, distance))
        results.append((path, float(distance)))  # 转换为标准float类型
    conn.close()
    results.sort(key=lambda x: x[1])
    return results[:top_k]

# 文件夹监控
class ImageHandler(FileSystemEventHandler):
    def on_created(self, event):
        if event.is_directory:
            return
        features = extract_features(event.src_path)
        insert_image(event.src_path, features)

observer = Observer()
event_handler = ImageHandler()
observer.schedule(event_handler, path=UPLOAD_FOLDER, recursive=False)

# API
app = Flask(__name__)


# 上传文件接口
@app.route('/upload', methods=['POST'])
def upload():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    # 保存文件到upload_folder目录
    file_path = os.path.join(UPLOAD_FOLDER, file.filename)
    file.save(file_path)
    return jsonify({"message": "File uploaded successfully", "file_path": file_path}), 200


# 批量上传文件接口
@app.route('/uploads', methods=['POST'])
def uploads():
    if 'files[]' not in request.files:
        return jsonify({"error": "No file part"}), 400

    files = request.files.getlist('files[]')
    if len(files) == 0:
        return jsonify({"error": "No selected files"}), 400

    file_paths = []
    for file in files:
        if file.filename == '':
            continue
        # 保存文件到upload_folder目录
        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
        file.save(file_path)
        file_paths.append(file_path)

    return jsonify({"message": "Files uploaded successfully", "file_paths": file_paths}),

# 搜索相似图片接口
@app.route('/search', methods=['POST'])
def search():
    file = request.files['image']
    file_path = os.path.join('temp', file.filename)
    file.save(file_path)
    query_feature = extract_features(file_path)
    results = search_similar_images(query_feature)
    # return jsonify(results)
    return jsonify(results)

# 提供访问上传文件的接口
@app.route('/images/<filename>')
def uploaded_file(filename):
    return send_from_directory(UPLOAD_FOLDER, filename)

if __name__ == '__main__':
    # 启动时初始化数据库和表
    init_db()
    # 开始监控文件夹
    observer.start()
    app.run(debug=True)
